from graphviz import Digraph


def draw_dfa(states, alphabet, transitions, start_state, accept_states, filename='just_pic'):
    """
    Draw a DFA using Graphviz.

    :param states: List of states in the DFA.
    :param alphabet: List of alphabet symbols.
    :param transitions: Dictionary where the keys are tuples (current_state, input_symbol) and
                        values are the next state.
    :param start_state: The starting state.
    :param accept_states: List of accepting states.
    :param filename: The name of the output file (without extension).
    """

    # Initialize the Digraph object
    dot = Digraph(format='png')

    # Set the graph attributes
    dot.attr(rankdir='LR', size='8,5')
    dot.graph_attr['dpi'] = '1000'
    # Add invisible start arrow
    dot.node('', shape='none', width='0')
    dot.edge('', start_state, label='Start')

    # Add states to the graph
    for state in states:
        if state in accept_states:
            dot.node(state, shape='doublecircle')
        else:
            dot.node(state, shape='circle')

    # Add transitions to the graph
    for (current_state, input_symbol), next_state in transitions.items():
        dot.edge(current_state, next_state, label=input_symbol)

    # Render the graph to a file
    dot.render(filename)
    print(f"DFA diagram saved as {filename}.png")


# Example usage
# states = ['S_0', 'S_1', 'S_2', 'S_3']
states = ['0','1', '2', '3']
# states = ['S_0', 'S_1', 'S_2','S_3', 'S_4', 'S_5','S_6', 'S_7', 'S_8']
alphabet = ['a', 'b']
transitions = {
    # ('S_0', 'a'): 'S_1',
    # ('S_0', 'b'): 'S_2',
    # ('S_1', 'a'): 'S_1',
    # ('S_1', 'b'): 'S_3',
    # ('S_2', 'a'): 'S_1',
    # ('S_2', 'b'): 'S_2',
    # ('S_3', 'a'): 'S_1',
    # ('S_3', 'b'): 'S_4',
    # ('S_4', 'a'): 'S_5',
    # ('S_4', 'b'): 'S_6',
    # ('S_5', 'a'): 'S_5',
    # ('S_5', 'b'): 'S_7',
    # ('S_6', 'a'): 'S_5',
    # ('S_6', 'b'): 'S_6',
    # ('S_7', 'a'): 'S_5',
    # ('S_7', 'b'): 'S_8',
    # ('S_8', 'a'): 'S_5',
    # ('S_8', 'b'): 'S_6',
    # ('S_3', 'a'): 'S_2',
    # ('S_3', 'b'): 'S_3',
    # ('0', 'a,b'): '0',
    # ('0', 'b'): '0',
    ('0', 'b'): '0',
    ('0', 'a'): '1',
    ('1', 'a'): '1',
    ('1', 'b'): '2',
    ('2', 'a'): '1',
    ('2', 'b'): '3',
    ('3', 'a,b'): '3',
    # ('3', 'b'): '3',


}
# start_state = 'S_0'
start_state = '0'
# accept_states = ['S_4', 'S_5', 'S_6', 'S_7', 'S_8']
# accept_states = ['0']
accept_states = ['3']

# Call the function to draw the DFA
draw_dfa(states, alphabet, transitions, start_state, accept_states)
